import cv2
import os
from cv2 import mean
import numpy as np
import sys
sys.path.append(os.path.abspath('.'))
from utils.util import sci_SSIM as SSIM 
from utils.util import SSIM_mean_param,SSIM_std_param,SSIM_cov_param
from .base_util import BaseUtil

def heat_map(ori_img,dst_img):
    heat = np.abs(dst_img - ori_img).astype(np.uint8)
    return heat


def corr_heat_map(ori_img,dst_img,bh,bw):

    H,W = dst_img.shape
    corr_mat = np.zeros((H,W))
    
    for h in range(1,int(H/bh)+1):
        for w in range(1,int(W/bw)+1):

            fake_vec = dst_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw ].reshape(bh*bw)
            real_vec = ori_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw ].reshape(bh*bw)
            corr_mat[(h-1)*bh : h*bh , (w-1)*bw : w*bw] = np.corrcoef(fake_vec,real_vec)[0,1]
            pass
    # corr_mat = ((corr_mat+1)*255/2).astype(np.uint8)
    corr_mat = (np.abs(corr_mat)*255).astype(np.uint8)
    # heat_corr = cv2.applyColorMap(corr_mat,cv2.COLORMAP_JET)
    return corr_mat

def SSIM_heat_map(ori_img,dst_img,bh,bw):
    
    H,W = dst_img.shape
    SSIM_mat = np.zeros((H,W))
    
    for h in range(1,int(H/bh)+1):
        for w in range(1,int(W/bw)+1):
            fake_vec = dst_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw ]
            real_vec = ori_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw]
            SSIM_mat[(h-1)*bh : h*bh , (w-1)*bw : w*bw] = SSIM(fake_vec,real_vec)

            pass
    SSIM_mat = (SSIM_mat*255).astype(np.uint8)
    # heat_ssim = cv2.applyColorMap(SSIM_mat,cv2.COLORMAP_JET)
    return SSIM_mat

def heatColor(img):
    return cv2.applyColorMap(img,cv2.COLORMAP_JET)

class CycleGANS012DFFTUtil(BaseUtil):
    def __init__(self, opt):
        BaseUtil.__init__(self,opt)
        
    def phaseData(self):
        record = []
        self.img = {}
        if not os.path.exists(self.opt.dir):
            return False
        self.opt.files = os.listdir(self.opt.dir)
        for file in self.opt.files:  
            tmp = file.split('.png_s0_')
            if 'fake' in tmp[1]:
                if not tmp[0] in record:
                    record.append(tmp[0])
                    fake_S01 = cv2.imread(os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S01.png'))
                    real_S01 = cv2.imread(os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S1.png'))
                    
                    fake_S02 = cv2.imread(os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S02.png'))
                    real_S02 = cv2.imread(os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S2.png'))
                    
                    fake_S0D = cv2.imread(os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S0D.png'))
                    real_S0D = cv2.imread(os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SD.png'))
                    
                    self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S01.png')] = cv2.cvtColor(fake_S01,cv2.COLOR_RGB2GRAY)
                    self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S1.png')] = cv2.cvtColor(real_S01,cv2.COLOR_RGB2GRAY)
                    
                    self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S02.png')] = cv2.cvtColor(fake_S02,cv2.COLOR_RGB2GRAY)
                    self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S2.png')] = cv2.cvtColor(real_S02,cv2.COLOR_RGB2GRAY)
                    
                    self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S0D.png')] = cv2.cvtColor(fake_S0D,cv2.COLOR_RGB2GRAY)
                    self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SD.png')] = cv2.cvtColor(real_S0D,cv2.COLOR_RGB2GRAY)
        return True
    
    def process(self):
        self.logfile = os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'uilt_log.txt')
        if os.path.exists(self.logfile):
            os.remove(self.logfile)
        for process in self.opt.processes:
            getattr(self,process)()
        # self.phaseData()
    def SSIM_tool(self):
        _dir = self.opt.dir
        record = []
        bh=self.opt.bh
        bw=self.opt.bw
        
        mean_S01_p1 = 0
        mean_S01_p2 = 0
        mean_S01_p3 = 0
        mean_S01_p123 = 0
        mean_S01_ssim = 0
        
        mean_S02_p1 = 0
        mean_S02_p2 = 0
        mean_S02_p3 = 0
        mean_S02_p123 = 0
        mean_S02_ssim = 0
        
        mean_S0D_p1 = 0
        mean_S0D_p2 = 0
        mean_S0D_p3 = 0
        mean_S0D_p123 = 0
        mean_S0D_ssim = 0
        
        count = 0
        print('SSIM tool start')
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if 'fake_' in tmp[1]:
                if not tmp[0] in record:
                    record.append(tmp[0])
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S01.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S1.png')]
                    # H,W = fake.shape
                    ssim = SSIM(fake,real)
                    p1 = SSIM_mean_param(real,fake,)
                    p2 = SSIM_std_param(real,fake)
                    p3 = SSIM_cov_param(real,fake)

                    mean_S01_p1 += p1
                    mean_S01_p2 += p2
                    mean_S01_p3 += p3
                    mean_S01_p123 += p1*p2*p3
                    mean_S01_ssim += ssim
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S02.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S2.png')]
                    ssim = SSIM(fake,real)
                    p1 = SSIM_mean_param(real,fake)
                    p2 = SSIM_std_param(real,fake)
                    p3 = SSIM_cov_param(real,fake)

                    mean_S02_p1 += p1
                    mean_S02_p2 += p2
                    mean_S02_p3 += p3
                    mean_S02_p123 += p1*p2*p3
                    mean_S02_ssim += ssim
                    
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S0D.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SD.png')]
                    ssim = SSIM(fake,real)
                    p1 = SSIM_mean_param(real,fake)
                    p2 = SSIM_std_param(real,fake)
                    p3 = SSIM_cov_param(real,fake)

                    mean_S0D_p1 += p1
                    mean_S0D_p2 += p2
                    mean_S0D_p3 += p3
                    mean_S0D_p123 += p1*p2*p3
                    mean_S0D_ssim += ssim
                    count+=1
                    pass
        mean_S01_p1 /= count
        mean_S01_p2 /= count
        mean_S01_p3 /= count
        mean_S01_p123 /= count
        mean_S01_ssim /= count
        mean_S02_p1 /= count
        mean_S02_p2 /= count
        mean_S02_p3 /= count
        mean_S02_p123 /= count
        mean_S02_ssim /= count
        mean_S0D_p1 /= count
        mean_S0D_p2 /= count
        mean_S0D_p3 /= count
        mean_S0D_p123 /= count
        mean_S0D_ssim /= count
        # logfile = os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'uilt_log.txt')
        x='None'
        with open(self.logfile,'a') as f:
            x = "S01: mean_SSIM: "+str(mean_S01_ssim)+" ,p1: "+str(mean_S01_p1)+" ,p2: "+str(mean_S01_p2)+" ,p3: "+str(mean_S01_p3)+" ,p123: "+str(mean_S01_p123)+'\n'
            x += "S02: mean_SSIM: "+str(mean_S02_ssim)+" ,p1: "+str(mean_S02_p1)+" ,p2: "+str(mean_S02_p2)+" ,p3: "+str(mean_S02_p3)+" ,p123: "+str(mean_S02_p123)+'\n'
            x += "S0D: mean_SSIM: "+str(mean_S0D_ssim)+" ,p1: "+str(mean_S0D_p1)+" ,p2: "+str(mean_S0D_p2)+" ,p3: "+str(mean_S0D_p3)+" ,p123: "+str(mean_S0D_p123)+'\n'
            f.write(x)
        print(x)
        # print("S01: mean_SSIM: "+str(mean_S01_ssim)+" ,p1: "+str(mean_S01_p1)+" ,p2: "+str(mean_S01_p2)+" ,p3: "+str(mean_S01_p3)+" ,p123: "+str(mean_S01_p123))
        # print("S02: mean_SSIM: "+str(mean_S02_ssim)+" ,p1: "+str(mean_S02_p1)+" ,p2: "+str(mean_S02_p2)+" ,p3: "+str(mean_S02_p3)+" ,p123: "+str(mean_S02_p123))
        # print("S0D: mean_SSIM: "+str(mean_S0D_ssim)+" ,p1: "+str(mean_S0D_p1)+" ,p2: "+str(mean_S0D_p2)+" ,p3: "+str(mean_S0D_p3)+" ,p123: "+str(mean_S0D_p123))
        pass
    def mean_heat_map(self):
        _dir = self.opt.dir
        
        # files = os.listdir(_dir)

        record = []
        mean_heat_mat_S01 = np.zeros((256,256)).astype(np.float)
        mean_heat_mat_S02 = np.zeros((256,256)).astype(np.float)
        mean_heat_mat_S0D = np.zeros((256,256)).astype(np.float)
        count = 0
        print('Heat start')
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if 'fake_' in tmp[1]:
                if not tmp[0] in record:
                    record.append(tmp[0])
                    #
                    # fake = cv2.imread(os.path.join(_dir,file))
                    # real = cv2.imread(os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png'))
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S01.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S1.png')]
                    # H,W = fake.shape
                    heat_mat = heat_map(real,fake)
                    mean_heat_mat_S01  = (heat_mat  + mean_heat_mat_S01*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_heat__S01.png'),heatColor(heat_mat))
                    
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S02.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S2.png')]
                    heat_mat = heat_map(real,fake)
                    mean_heat_mat_S02  = (heat_mat  + mean_heat_mat_S02*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_heat_S02.png'),heatColor(heat_mat))
                    
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S0D.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SD.png')]
                    heat_mat = heat_map(real,fake)
                    mean_heat_mat_S0D  = (heat_mat  + mean_heat_mat_S0D*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_hcorr_S0D.png'),heatColor(heat_mat))
                    
                    count += 1
                    pass
                
            # img = cv2.imread(os.path.join(_dir,file))
        # mean_heat_mat_S01 = mean_heat_mat_S01
        self.mean_heat_S01_score = np.sum(mean_heat_mat_S01) / (mean_heat_mat_S01.shape[0]*mean_heat_mat_S01.shape[1])
        # self.mean_corr_score = np.where(mean_corr_mat>mean_corr_mat.mean())[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_heat_S01.png'),heatColor(mean_heat_mat_S01.astype(np.uint8)))
        # mean_heat_mat_S02 = mean_heat_mat_S02
        self.mean_heat_S02_score = np.sum(mean_heat_mat_S02) / (mean_heat_mat_S02.shape[0]*mean_heat_mat_S02.shape[1])
        # self.mean_corr_score = np.where(mean_corr_mat>mean_corr_mat.mean())[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_heat_S02.png'),heatColor(mean_heat_mat_S02.astype(np.uint8)))
        # mean_heat_mat_S0D = mean_heat_mat_S0D
        self.mean_heat_S0D_score = np.sum(mean_heat_mat_S0D) / (mean_heat_mat_S0D.shape[0]*mean_heat_mat_S0D.shape[1])
        # self.mean_corr_score = np.where(mean_corr_mat>mean_corr_mat.mean())[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_heat_S0D.png'),heatColor(mean_heat_mat_S0D.astype(np.uint8)))
        
        # print('S1  heat score:'+str(self.mean_heat_S01_score))
        # print('S2  heat score:'+str(self.mean_heat_S02_score))
        # print('SD  heat score:'+str(self.mean_heat_S0D_score))
        
        
        x='None'
        with open(self.logfile,'a') as f:
            x = 'S01  heat score:'+str(self.mean_heat_S01_score)+'\n'
            x += 'S02  heat score:'+str(self.mean_heat_S02_score)+'\n'
            x += 'S0D  heat score:'+str(self.mean_heat_S0D_score)+'\n'
            f.write(x)
        print(x)
        return mean_heat_mat_S01,mean_heat_mat_S02,mean_heat_mat_S0D

    def mean_corr_heat_map(self):
        _dir = self.opt.dir
        bh=self.opt.bh
        bw=self.opt.bw
        
        # files = os.listdir(_dir)

        record = []
        mean_corr_mat_S01 = np.zeros((256,256)).astype(np.float)
        mean_corr_mat_S02 = np.zeros((256,256)).astype(np.float)
        mean_corr_mat_S0D = np.zeros((256,256)).astype(np.float)
        count = 0
        H,W = self.opt.imgh, self.opt.imgw
        print('corr start')
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if 'fake_' in tmp[1]:
                if not tmp[0] in record:
                    record.append(tmp[0])
                    #
                    # fake = cv2.imread(os.path.join(_dir,file))
                    # real = cv2.imread(os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png'))
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S01.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S1.png')]
                    # H,W = fake.shape
                    corr_mat = corr_heat_map(real,fake,bh,bw)
                    mean_corr_mat_S01  = (corr_mat  + mean_corr_mat_S01*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S01.png'),heatColor(corr_mat))
                    
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S02.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S2.png')]
                    corr_mat = corr_heat_map(real,fake,bh,bw)
                    mean_corr_mat_S02  = (corr_mat  + mean_corr_mat_S02*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S02.png'),heatColor(corr_mat))
                    
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S0D.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SD.png')]
                    corr_mat = corr_heat_map(real,fake,bh,bw)
                    mean_corr_mat_S0D  = (corr_mat  + mean_corr_mat_S0D*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S0D.png'),heatColor(corr_mat))
                    
                    count += 1
                    pass
                
            # img = cv2.imread(os.path.join(_dir,file))
        mean_corr_mat_S01 = mean_corr_mat_S01
        self.mean_corr_S01_score = np.where(mean_corr_mat_S01>255*0.6)[0].shape[0]/(H*W)
        # self.mean_corr_score = np.where(mean_corr_mat>mean_corr_mat.mean())[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S01.png'),heatColor(mean_corr_mat_S01.astype(np.uint8)))
        
        mean_corr_mat_S02 = mean_corr_mat_S02
        self.mean_corr_S02_score = np.where(mean_corr_mat_S02>255*0.6)[0].shape[0]/(H*W)
        # self.mean_corr_score = np.where(mean_corr_mat>mean_corr_mat.mean())[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S02.png'),heatColor(mean_corr_mat_S02.astype(np.uint8)))
        
        mean_corr_mat_S0D = mean_corr_mat_S0D
        self.mean_corr_score = np.where(mean_corr_mat_S0D>255*0.6)[0].shape[0]/(H*W)
        # self.mean_corr_score = np.where(mean_corr_mat>mean_corr_mat.mean())[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S0D.png'),heatColor(mean_corr_mat_S0D.astype(np.uint8)))
        
        
        # logfile = os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'uilt_log.txt')
        x='None'
        with open(self.logfile,'a') as f:
            x='S01  corr score:'+str(self.mean_corr_S01_score)+'\n'
            x+='S02 corr score:'+str(self.mean_corr_S02_score)+'\n'
            x+='S0D corr score:'+str(self.mean_corr_score)+'\n'
            f.write(x)
        print(x)
            
        # print('S1  corr score:'+str(self.mean_corr_S01_score))
        # print('S2 corr score:'+str(self.mean_corr_S02_score))
        # print('SD corr score:'+str(self.mean_corr_score))    
        return mean_corr_mat_S01,mean_corr_mat_S02,mean_corr_mat_S0D

    def mean_SSIM_heat_map(self):
        _dir = self.opt.dir
        bh=self.opt.bh
        bw=self.opt.bw
        
        files = os.listdir(_dir)
                    
        record = []
        mean_SSIM_S01 = np.zeros((256,256)).astype(np.float)
        mean_SSIM_S02 = np.zeros((256,256)).astype(np.float)
        mean_SSIM_S0D = np.zeros((256,256)).astype(np.float)
        count = 0
        H,W = self.opt.imgh, self.opt.imgw
        print('SSIM start')
        for file in files:
            tmp = file.split('.png_s0_')
            if 'fake_' in tmp[1]:
                if not tmp[0] in record:
                    record.append(tmp[0])
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S01.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S1.png')]
                    # H,W,_ = fake.shape
                    SSIM_mat = SSIM_heat_map(real,fake,bh,bw)
                    mean_SSIM_S01 = (SSIM_mat + mean_SSIM_S01*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S01.png'),heatColor(SSIM_mat))
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S02.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S2.png')]
                    # H,W,_ = fake.shape
                    SSIM_mat = SSIM_heat_map(real,fake,bh,bw)
                    mean_SSIM_S02 = (SSIM_mat + mean_SSIM_S02*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S02.png'),heatColor(SSIM_mat))
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S0D.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SD.png')]
                    # H,W,_ = fake.shape
                    SSIM_mat = SSIM_heat_map(real,fake,bh,bw)
                    mean_SSIM_S0D = (SSIM_mat + mean_SSIM_S0D*count)/(count+1)
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S0D.png'),heatColor(SSIM_mat))
                    
                    count += 1
                    pass
                
        # self.mean_SSIM_score = np.where(np.unique(mean_SSIM_mat)>255*0.7)[0].shape[0]/np.unique(mean_SSIM_mat).shape[0]?
        self.mean_SSIM_score_S01 = np.where(mean_SSIM_S01>255*0.6)[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S01.png'),heatColor(mean_SSIM_S01.astype(np.uint8)))
        
        self.mean_SSIM_score_S02 = np.where(mean_SSIM_S02>255*0.6)[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S02.png'),heatColor(mean_SSIM_S02.astype(np.uint8)))
        
        self.mean_SSIM_score_S0D = np.where(mean_SSIM_S0D>255*0.6)[0].shape[0]/(H*W)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'_S0D.png'),heatColor(mean_SSIM_S0D.astype(np.uint8)))
        
        # logfile = os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'uilt_log.txt')
        x='None'
        with open(self.logfile,'a') as f:
            x='S01  SSIM score:'+str(self.mean_SSIM_score_S01)+'\n'
            x+='S02 SSIM score:'+str(self.mean_SSIM_score_S02)+'\n'
            x+='S0D SSIM score:'+str(self.mean_SSIM_score_S0D)+'\n'
            f.write(x)
        print(x)
        
        # print('S01 SSIM score:'+str(self.mean_SSIM_score_S01))
        # print('S02 SSIM score:'+str(self.mean_SSIM_score_S02))
        # print('S0D SSIM score:'+str(self.mean_SSIM_score_S0D))
        return mean_SSIM_S01,mean_SSIM_S02,mean_SSIM_S0D
    def FFID(self):
        _dir = self.opt.dir
        files = os.listdir(_dir)
        record = []
        from  numpy.fft import fft2,fftshift
        
        print('FFID start')
        for file in files:
            tmp = file.split('.png_s0_')
            if 'fake_' in tmp[1]:
                if not tmp[0] in record:
                    record.append(tmp[0])
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S01.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S1.png')]
                    F_fake = fftshift(fft2(fake))
                    F_real = fftshift(fft2(real))
                    FFID_S01 = np.sum(F_fake-F_real)
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S02.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_S2.png')]
                    F_fake = fftshift(fft2(fake))
                    F_real = fftshift(fft2(real))
                    FFID_S02 = np.sum(F_fake-F_real)
                    
                    fake = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_fake_S0D.png')]
                    real = self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SD.png')]
                    F_fake = fftshift(fft2(fake))
                    F_real = fftshift(fft2(real))
                    FFID_S0D = np.sum(F_fake-F_real)
        pass
    

